package com.cfh.springcloudservletgateway.loadbalance;

import com.netflix.loadbalancer.BaseLoadBalancer;
import com.netflix.loadbalancer.Server;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class ZookeeperLoadBalancer extends BaseLoadBalancer {

    private final DiscoveryClient discoveryClient;

    private Map<String, BaseLoadBalancer> loadBalancerMap = new ConcurrentHashMap<>();

    @Autowired
    public ZookeeperLoadBalancer(DiscoveryClient discoveryClient) {
        this.discoveryClient = discoveryClient;
        updateServer();
    }

    @Scheduled(fixedRate = 10 * 1000)
    private void updateServer() {

        discoveryClient.getServices().stream().forEach(serviceName -> {

            BaseLoadBalancer baseLoadBalancer = new BaseLoadBalancer();
            loadBalancerMap.put(serviceName, baseLoadBalancer);

            discoveryClient.getInstances(serviceName).stream().forEach(serviceInstance -> {
                Server server = new Server(serviceInstance.isSecure() ? "https://" : "http://",
                        serviceInstance.getHost(), serviceInstance.getPort());
                baseLoadBalancer.addServer(server);
            });

        });
    }

    @Override
    public Server chooseServer(Object key) {
        if (key instanceof String) {
            String serviceName = String.valueOf(key);
            BaseLoadBalancer baseLoadBalancer = loadBalancerMap.get(serviceName);
            List<Server> allServers = baseLoadBalancer.getAllServers();
            return allServers.get(0);
        }
        return super.chooseServer(key);
    }


}
